#!/usr/bin/python

from openpyxl.reader.excel import load_workbook

#获取混合样本的数据
def get_sample_data(filename):
    #保存样本数据的三维字典{样本号：{位点：{等位基因：峰高}}
    samples_data = {}
    f = open(filename,'r')
    while True:
        line = f.readline().rstrip()
        if not line:
            break
        # lst = [a for a in line.split('\t') if a != '']
        if not line.startswith('Sample'):
            lst = []
            for a in line.split('\t'):
                if a != '':
                    lst.append(a)
            alle_and_height = lst[3:]
            allle_cnt = int(len(alle_and_height)/2)
            alleles = alle_and_height[:allle_cnt]
            heights = alle_and_height[allle_cnt:]
            sample_name = (lst[0],lst[1])
            loci_name = lst[2]
            for allele,height in zip(alleles,heights):
                if sample_name not in samples_data:
                    samples_data[sample_name] = {loci_name: {allele:height}}
                else:
                    if loci_name not in samples_data[sample_name]:
                        samples_data[sample_name][loci_name] = {allele:height}
                    else:
                        samples_data[sample_name][loci_name][allele] = height
    f.close()
    return samples_data

#读取男女样本的分型
def get_sample_genotype(filename):
    male_genotype = {}
    female_genotype = {}
    wb = load_workbook(filename,read_only=True)
    ws = wb.active
    for i in range(2,ws.max_row+1):
        loci = ws.cell(row=i, column=1).value
        male = str(ws.cell(row=i, column=3).value)
        female = str(ws.cell(row=i, column=2).value)
        male_alleles = [a for a in male.split(',') if a]
        female_alleles = [a for a in female.split(',') if a]
        male_genotype[loci] = male_alleles
        female_genotype[loci] = female_alleles
    wb.close()
    return male_genotype, female_genotype

#读取指定分型的样本的峰高
def get_sample_height(genotype,sample_data):
    sample_height = {}
    for loci in genotype:
        height_lst = []
        if loci in sample_data:
            for allele in genotype[loci]:
                if allele in sample_data[loci]:
                    height_lst.append(int(sample_data[loci][allele]))
        height = 0
        if len(height_lst) == 1:
            height = height_lst[0]/2
        if len(height_lst) == 2:
            height = height_lst[0]/2 + height_lst[1]/2
        sample_height[loci] = height
    return sample_height

#求比例
def get_ratio_of_female_and_male(male_height,female_height):
    ratios = {}
    for loci in male_height:
        if loci in female_height:
            ratio = female_height[loci]/male_height[loci]
            ratios[loci] = ratio
    return ratios

if __name__ == '__main__':
    filename = '../testData/20250820-casework_SamplePlotGenotypeTable.txt'
    samples_data = get_sample_data(filename)
    filename = '../testData/男女样本分型.xlsx'
    male_genotype, female_genotype = get_sample_genotype(filename)

    for sample in samples_data:
        hasMale = True
        male_height = get_sample_height(male_genotype,samples_data[sample])
        female_height = get_sample_height(female_genotype,samples_data[sample])
        for loci in male_height:
            if male_height[loci] == 0:
                hasMale = False
        if hasMale:
            ratios = get_ratio_of_female_and_male(male_height,female_height)
            print (f'样本名称:{sample}')
            print ('基因座\t女性峰高\t男性峰高\t比例')
            for loci in ratios:
                print (loci,female_height[loci],male_height[loci],'%.2f' % ratios[loci],sep='\t')